from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier  # 引入决策树分类器
from sklearn.metrics import accuracy_score  # 计算分类准确率

if __name__ == '__main__':
    wine = load_wine()
    # 划分训练集与测试集
    features_train, features_test, labels_train, labels_test = \
        train_test_split(wine.data, wine.target, test_size=0.2, random_state=80)
    decision_tree_classifier = DecisionTreeClassifier(criterion="entropy")  # 以信息增益作为分类依据
    decision_tree_classifier.fit(features_train, labels_train)
    # 计算训练集与测试集的预测准确率 (通过 accuracy_score 与 decision_tree_classifier.score 两种方式)
    accuracy_train = decision_tree_classifier.score(features_train, labels_train)
    labels_test_predict = decision_tree_classifier.predict(features_test)
    accuracy_test = accuracy_score(labels_test, labels_test_predict)
    print("训练集准确率:", accuracy_train)
    print("测试集准确率:", accuracy_test)
